#!/usr/bin/env python

import ssummo.ssummolib as slib
import cPickle as pickle
import colours
import re
#from ssummo import colours
import threading, Queue
import multiprocessing
import random
#import numpy as np
import sys, os.path
from rankAbundance import Indices

try:
    import matplotlib.pyplot as plt
    matplot = True
except ImportError:
    matplot = False
    sys.stderr.write("matplotlib import error. Graph-plotting functionality will not work. ")

class LocalOptions( slib.Options ):
    def __init__(self,args=None):
        slib.Options.__init__(self)
        self.singleargs = ['--draw','-replicates','-increment','-collapse-at-rank']
        self.multiargs = ['-in','-out','-order','-indices']
        self.helpTxt = { '-in' : 'Input files.',
                         '-out': 'Output files. Default: standard out',
                         '-replicates': 'Number of replicates for each dataset. [5]',
                         '-increment' : 'For rarefaction curves, only sample every `-increment` species. [100]',
                         '--draw':'If results are already saved, just draw the results saved in file `-in`',
                         '-collapse-at-rank': 'Collapse results at specific rank.',
                         '--indices'  :  'Calculate and plot ALL biodiversity indices.',
                         '-indices'   :  ' List which indices to plot [shannon | shannon_max | simpson | jack].',
                         '-order' : 'Defines the order when drawing old results.',
                         }
        self.options.update( { '-replicates' : None    ,
                               '-increment'  : None    ,
                               '--draw'      : False   ,
                               '-collapse-at-rank': None,
                               '--indices'   : False   ,
                               '-indices'    : []      ,
                               '-out' : []             ,
                             } )
        self.groups = None
        self.switches = {'--draw' : False ,
                         '--indices' : False}
        self.post_checks += [ self.groupify , self.indices]
        if args is not None:
            self.parseArgs(args)
    def groupify( self ):
        """If there are any groups defined on the command line,
        this'll just create a list contianing the lengths of 
        each group.
        """
        if len( self.regopts ) > 0:
            self.groups = [ len( in_files ) for in_files in [ self.regopts[key] for key in sorted(self.regopts.keys()) ]]
    def indices( self ):
        """If you specify --indices, append all."""
        if self.options['--indices']:
            self.options['-indices'] = [ 'shannon','shannon_max','simpson','jack',]
        elif len( self.options['-indices'] ) > 0:
            self.options['--indices'] = True

def get_accession_places( node, accessions={} , path='/' ):
    for key in node.keys():
        if key == 'accessions':
            for accession in node[key]:
                accessions.update( { accession : path } )
        else:
            accessions = get_accession_places( node[key] , accessions=accessions , path=os.path.join(path,key) )
    return accessions

## These two functions from:-
## http://matplotlib.sourceforge.net/examples/pylab_examples/multiple_yaxis_with_spines.html 
def make_patch_spines_invisible(ax):
    ax.set_frame_on(True)
    ax.patch.set_visible(False)
    for sp in ax.spines.itervalues():
        sp.set_visible(False)

def make_spine_invisible(ax, direction):
    if direction in ["right", "left"]:
        ax.yaxis.set_ticks_position(direction)
        ax.yaxis.set_label_position(direction)
        anti = 'right' if direction == 'left' else 'left'
    elif direction in ["top", "bottom"]:
        ax.xaxis.set_ticks_position(direction)
        ax.xaxis.set_label_position(direction)
        anti = 'bottom' if direction == 'top' else 'top'
    else:
        raise ValueError("Unknown Direction : %s" % (direction,))
    ax.spines[direction].set_visible(True)
    ax.spines[anti].set_visible( False )

class PlotRarefaction( ):
    def __init__(self,options): # datasets,replicates=None,increment=None , groups=None):
        """
        """
        if options['-replicates'] is None:
            self.replicates = 5
        else:
            self.replicates = int(options['-replicates'])
        if options['-increment'] is None:
            self.increment = 1000
        else:
            self.increment = int(options['-increment'])

        self.col_ind = 0 # For assigning colours.
        number_datasets =  len( options['-in'] ) 
        self.n = number_datasets
        self.dataset_names = []
        self.options = options
        if '_groups' in self.options.options:
            self.colours = [colours.HSL_to_HEX(colour) for colour in colours.Heat( N=self.n , groups=self.options['_groups'] ) ]
        else:
            self.colours = [colours.HSL_to_HEX(colour) for colour in colours.Heat( N=self.n ) ]

        if matplot:
            self.fig = plt.figure()
            self.axis = self.fig.add_subplot( 111 )
            plt.ioff()  # Turn off interactive mode (default anyway, unless changed in matplotlib config file); speeds things up.
            plt.xlabel( 'Number of randomly chosen sequences' )
            plt.ylabel( 'Number of distinct taxa' )
            if self.options['--indices']:
                self.init_axes()
        self.markers = [ 'x','+','o','v','d','<','s','8' ]
        self.marker_sizes = dict( [(key, 49)  for key in self.markers ] )
        self.marker_sizes['o'] = 9 
        self.marker_ind = 0

        self.artists = []
        self.index_artists = []
        self.index_heights = []
        self.index_labels = []

        self.x_max = 0
        self.y_max = 0

    def reinit( self ):
        """To be called if need to reload the colours. This needs doing with
        --draw, when PlotRarefaction is initiated before the previous results
        have been loaded (i.e. before the original options have been loaded).
        """
        self.n = len( self.options['-in'] )
        print self.options['_groups']
        self.colours = [colours.HSL_to_HEX(colour) for colour in colours.Heat( N=self.n , groups=self.options['_groups'] ) ]


    def init_axes( self ):
        """Here define methods to initialise additional axes. This is called
        during initialisation if options['--indices'] is True.
        """
        axis_adjustment = 0.05
        setattr( self , 'fig2' , plt.figure() )
        setattr( self, 'axis2' , self.fig2.add_subplot( 111 ) )
        self.axis2.set_xlabel( 'Number sequences sampled' )
        cur_axis = self.axis2
        left_increment = 0.
        right_increment = 1.
        axis_ind = 3
        n_axes = len( self.options['-indices'] )
        if n_axes > 1:
            side = 'right'
            for counter in xrange( n_axes ) :
                # Move the y-axes across.
                if side == 'left':
                    new_axis = self.fig2.add_subplot( 111 ,axisbg='#000000')
                    left_increment -= axis_adjustment
                    new_axis.spines[side].set_position( ('axes' ,left_increment ) )
                    make_patch_spines_invisible( new_axis )
                    make_spine_invisible( new_axis ,side)
                    side = 'right'
                elif side == 'right':
                    new_axis = cur_axis.twinx()
                    #if axis_ind > 3:
                    new_axis.spines[side].set_position( ('axes' ,right_increment ) )
                    make_patch_spines_invisible( new_axis )
                    make_spine_invisible( new_axis ,side)
                    #else:
                    #    new_axis.spines['left'].set_position( ('axes',right_increment ) )
                    #make_patch_spines_invisible( new_axis )
                    #    make_spine_invisible( new_axis ,'left')
                    right_increment += axis_adjustment
                    side = 'left'
                else: raise IOError

                axis_id = 'axis{0}'.format( axis_ind )
                print 'creating ',axis_id
                setattr( self , axis_id , new_axis )
                axis_ind += 1

                cur_axis = new_axis

        self.fig2.subplots_adjust( left=0.15,right=0.85 )
        self.axes = range( 2 , axis_ind )
        self.cur_axis  = 2

    def plot_scatter( self, xvals , yvals, marker=None,label=None ,*args , **kwargs):
        """xvals is directly given to plt.scatter(), so must be compatible with 
        matplotlib.

        y_vals can optionally be a dictionary of vectors, in which case all the
        keys must be present in xvals. If not a dictionary, assume that it is a
        sequence with same length as xvals.
        
        marker can be given a number indexing which of self.markers to use, or
        a string to pass directly to matplotlib. By default, use whatever marker
        is selected with:
            self.markers[ self.marker_ind ]

        label can optionally be used to indicate the label in the legend.

        Optional *args and **kwargs are passed directly to plt.scatter()
        """
        all_xvals = [] ; all_yvals = []
        if marker is None:
            marker = self.markers[self.marker_ind] 
        elif type(marker) is int:
            marker = self.markers[marker]
        if type( yvals ) is dict:
            for n in xrange( len( yvals[xvals[0]] ) ):
                all_xvals.append( xvals )
                all_yvals.append( [ yvals[x][n] for x in xvals ] )
        else:
            all_xvals = xvals ; all_yvals = yvals
        cur_axis = getattr( self , 'axis{0}'.format(self.cur_axis) )
        # tests.
        try:
            mark_size = self.marker_sizes[marker]
            col = self.colours[self.col_ind]
        except IndexError:
            print '\n',len( self.marker_sizes ) , marker
            print len( self.colours ) , self.col_ind
            raise

        collection = cur_axis.scatter( all_xvals , all_yvals ,marker=marker,s=self.marker_sizes[marker],color=self.colours[self.col_ind], label=label, *args, **kwargs )
        return collection

    def plot_scatter_errors( self, xvals , yvals , mid=None, xerr=None, yerr=None, label=None, *args, **kwargs ):
        """Plot xvals against yvals with error bars.
        Uses self.plot_scatter to draw the points, but adds error bars.

        mid - This is the function which calculates the centre of crosses.
            By default, calculate the mean.

        xerr - By default, don't draw x error bars. Otherwise, pass a function
            reference which will return the error values when passed xvals.

        yerr - By default, draw 95% Confidence Intervals ( i.e. +/- 1.96 *  1STD / sqrt( replicates ) ).
        """
        if mid is None:
            mid = plt.np.mean
        mean_y_vals = [ mid(yval) for yval in yvals ]
        collection = self.plot_scatter( xvals, mean_y_vals, label=label,*args,**kwargs )
        xerrors = None
        if xerr is not None:
            xerrors = xerr( xvals )
        if yerr is None:
            yerrors =  [ 1.96 * ( plt.np.std( yval ) / plt.np.sqrt( self.replicates ) ) for yval in yvals ]
            yerrors =  [ plt.np.std( yval )  for yval in yvals ]
        else:
            yerrors = yerr( yvals )
        if label is not None:
            sys.stdout.write( '{0}\t'.format( label ) ) 
        sys.stdout.write( '\t'.join( [ '{0:.3f} +- {1:.3f}'.format( mean_y_vals[i] , yerrors[i] ) for i in xrange( len( yvals ) ) ]) )
        sys.stdout.write( '\n\n' )
        y2_max = max( mean_y_vals ) + max( yerrors )
        cur_axis = getattr( self , 'axis{0}'.format(self.cur_axis) )
        (plotline, caplines, barlinecols) = cur_axis.errorbar( xvals , mean_y_vals,c=self.colours[self.col_ind], xerr=xerrors, yerr=yerrors, linestyle='None',label=label,*args, **kwargs )
        plt.setp( barlinecols , linestyles='dotted' )
        if cur_axis.get_ylim()[1] < y2_max:
            cur_axis.set_ylim( 0 , int(y2_max) + 1 )
        return collection


    def plot_box( self, xvals , yvals , label=None , widths=None , *args, **kwargs ):
        """Given one or more x-value, will plot all given y-values
        as a box-and-whisker plot.
        This is designed to be called once per dataset, so that each dataset
        has the same colour.
        See http://matplotlib.sourceforge.net/api/pyplot_api.html#matplotlib.pyplot.boxplot
        for more information"""
        all_xvals = [] ; all_yvals = []
        if widths is None:
            widths = self.increment * 0.3
        mappings = self.axis.boxplot( yvals , positions=xvals, widths=widths,*args , **kwargs )
        for mapping in mappings:
            plt.setp( mappings[mapping] , color = self.colours[self.col_ind] )
        # mappings.keys() = ['medians', 'fliers', 'whiskers', 'boxes', 'caps']
        plt.setp( mappings['whiskers'] , linestyle='dotted' )
        self.col_ind += 1
        return mappings

    def draw( self ):
        """Just provide the x-axis and y-axis limits. The x-axis is rounded
        up to the nearest 100, and the y-axis is rounded up to the nearest
        500.
        """
        self.axis.legend( self.artists , self.options['-in'] , loc=4, borderaxespad=0.2 )
        if self.options['--indices']:
            self.draw_indices( self.index_artists, self.index_labels )

        self.axis.set_xlim( 0 , 500*(1+(self.x_max//500)) ) # round up to the nearest 500
        self.axis.set_ylim( 0 , 100*(1+(self.y_max//100)) ) # round up to the nearest 100
        self.ticker = plt.matplotlib.ticker.MultipleLocator( base = self.increment * 2)
        self.axis.xaxis.set_major_locator( self.ticker )
        self.subticker = plt.matplotlib.ticker.MultipleLocator( base = self.increment  )
        self.axis.xaxis.set_minor_locator( self.subticker )
        if self.options['--indices']:
            i = 2
            while hasattr( self , 'axis{0}'.format( i) ):
                ax = getattr( self, 'axis{0}'.format(i) )
                ax.xaxis.set_major_locator( self.ticker )
                ax.xaxis.set_minor_locator( self.subticker )
                i += 1
        self.axis.grid( True )
        plt.show()

    def draw_indices( self , index_artists, index_labels ):
        self.axis2.legend( index_artists , index_labels ,loc=4, borderaxespad=0.2 )
        index = 2
        for i, bio_index in enumerate( self.options['-indices'] ):
            axis_id = 'axis{0}'.format( index )
            if hasattr( self , axis_id ):
                print 'setting {0} on {1}'.format( bio_index , axis_id )
                cur_axis = getattr( self , axis_id )
                cur_axis.set_ylabel( bio_index )
                cur_axis.tick_params( axis='y',size=4,width=1.5 )
                index += 1
            else:
                print '{0} does not exist'.format( axis_id )

    def load_rarefaction( self, data ):
        # Give matplotlib the data to plot behind the scenes. 
        # We keep the 'artists' to colour later.
        #xvals = []; yvals = []
        #for xval in sorted( data.rarefaction.keys() ):
        #    xvals.append( xval ) ; yvals.append( data.rarefaction[xval] )

        labels = self.options['-in']
        if len( data.rarefaction ) == 0:
            return
        if self.options['--indices']:
            self.cur_axis = 2
            sys.stdout.write( '{0}\t'.format( data.file_name ) )
            # -indices = [ 'jack' , 'shannon' , 'shannon_max' , 'simpson' ]
            for bio_index in self.options['-indices']:
                if hasattr( data , bio_index ):
                    index_data = [ getattr( data , bio_index )[xval] for xval in data.xvals ]
                    #plotter.cur_axis = plotter.axes[ plotter.marker_ind ]
                    collection = self.plot_scatter_errors( data.xvals, index_data, label=bio_index )
                    self.marker_ind += 1
                    self.index_artists.append( collection )
                    self.index_labels.append( '{0} {1}'.format( data.file_name ,bio_index ) )
                    self.cur_axis += 1
            self.marker_ind = 0

        #self.x_max = max( [data.xvals , self.x_max ])
        self.x_max = max( data.xvals[-1] , self.x_max )
        yvals = [ data.rarefaction[xval] for xval in data.xvals ]
        self.y_max = max( yvals[-1] + [ self.y_max ])
        mappings = self.plot_box(data.xvals, yvals)
        self.artists.append( mappings['boxes'][0] )

    def resume( self , results , loc=4 ):
        """results should be a dictionary of Results() instances.
        If you want them presented in a specific order, then give 
        the keys of the results dictionary in the order which you
        want them displayed.
        loc is the location of the legend according to matplotlib 
        co-ordinates.
        """
        artists = []
        data_names = []
        max_x = 0
        max_y = 0
        order = [ truncate_name( infile ) for infile in self.options['-in'] ]
        try:
            iterator = [ results[key] for key in order ]
        except KeyError:
            print results.keys()
            print key
            raise
        self.reinit()

        if self.options['-increment'] is None:
            xvals = []
            for res in iterator:
                if len(res.xvals) > 0: xvals.append( max( res.xvals ) )
            increment = calc_increment( xvals )
            self.increment = increment

        for results_obj in iterator:
            self.load_rarefaction( results_obj )
#            print data_file, len( results_obj.xvals ) , len(results_obj.rarefaction )
#            xvals,yvals = results_obj.xvals , results_obj.rarefaction
#            max_x = max( max(xvals) , max_x )
#            max_y = max( max(yvals) , max_y )
#            data_names.append( data_file )
#            mappings = self.plot_box( xvals , yvals )
#            artists.append( mappings['boxes'][0] )
#        plt.legend( artists , data_names , loc=loc , borderaxespad=0.2 )
        self.draw( )
        pass

class Results( object ):
    __slots__ = [ 'file_name' , 'xvals' , 'rarefaction' , 'simpson' , 'shannon' ,'shannon_max' , 'jack' ]
    def __init__(self):
        """Class to store results.
        xvals is a list of all x-axis co-ordinates. This
        is (so far) always the number of sequences randomly
        sampled. Every other attribute - except for file_name,
        which is a string - is a dictionary, with
        each value in xvals being the keys, and values being
        lists containing Replicate values..
        """
        self.file_name = str()
        self.xvals = []
        self.rarefaction = {}
        self.simpson = {}
        self.shannon = {}
        self.shannon_max = {}
        self.jack = {}
    def __get__(self,val):
        return getattr( self , val )
    def __set__( self, key, val ):
        if type( val ) == Replicate:
            self.replicates[key] = val
        else:
            raise AttributeError( "{0} must be a Replicate instance".format( val.__name__ ) )
    def append( self, replicate ):
        if type( replicate ) != Replicate:
            raise AttributeError( "replicate must be a Replicate instance" )
        xval = replicate.nseqs
        if xval not in self.xvals:
            self.xvals.append( xval )
            self.rarefaction.update( { xval : [] } )
            self.simpson.update( { xval : [] } )
            self.shannon.update( { xval : [] } )
            self.shannon_max.update( { xval : [] } )
            self.jack.update( { xval : [] } )
        self.rarefaction[xval].append( replicate.rarefaction )
        self.simpson[xval].append( replicate.simpson )
        self.shannon[xval].append( replicate.shannon )
        self.shannon_max[xval].append( replicate.shannon_max )
        self.jack[xval].append( replicate.jack )

class Replicate( object ):
    __slots__ = [ 'nseqs', 'rarefaction' ,'shannon' , 'shannon_max' , 'simpson' , 'jack' ]
    def __init__(self):
        self.nseqs = int()
        for slot in self.__slots__:
            setattr( self , slot , float() )
    def __add__(self , other ):
        new = Replicate()
        for slot in self.__slots__:
            setattr( new, slot , getattr( self, slot ) + getattr( other , slot ) )
        return new
    def __div__(self,other):
        new = Replicate()
        for slot in self.__slots__:
            setattr( new , slot , getattr( self , slot ) / getattr( other , slot ) )
        return new



class Loader( multiprocessing.Process ):
    """
    We have one Loader instance for each SSuMMo results file. 
    When run is called, (usually with Loader().start() to run as
    a separate process), a Results instance will be created and 
    returned via self.queue, if it exists, or just returned
    otherwise
    """
    tax_index = slib.load_index()
    def __init__( self, results_dict, options, queue=None, semaphore=None ): #replicates=None,increment=None):
        """This can optionally be run as a separate process, in which case
        define a queue and a semaphore in the master process, and pass them
        to Loader when initialising. e.g.
        >>> sem = multiprocessing.Semaphore()
        >>> for index in xrange( number of processes ):
        >>>     ...
        >>>     instance = Loader( results , options, multiprocessing.Queue() ,
        >>>     semaphore= sem )
        >>>     sem.acquire()
        >>>     instance.start()
        
        Look at main() to see the full implementation.
        """
        if queue is not None:
            multiprocessing.Process.__init__(self)
        self.queue = queue
        self.semaphore = semaphore
        self.accession_places = get_accession_places( results_dict , accessions={} )
        self.nOTUs = len( set( self.accession_places.values() ) )
        self.nSeqs = len( self.accession_places )   # X-axis
        self.results_dict = results_dict
        self.options = options

        if options['-replicates'] is None:
            self.replicates = 5
        else:
            self.replicates = int(options['-replicates'])

        if options['-increment'] is None:
            self.increment = self.nSeqs / 20
        else:
            self.increment = int(options['-increment'])
        self.indices = None
        if options['--indices'] > 0:
            self.indices = True

    def _calc_indices( self, n_seqs_sampled, nodes ):
        """
        """
        counts = sorted( [ len(accessions) for accessions in nodes.values() ] ,reverse=True )
        probs = [ float(count) / n_seqs_sampled for count in counts ]                
        total = n_seqs_sampled
        biodiversity_calculator = Indices( self.options )
        replicate = Replicate()
        replicate.nseqs = total
        replicate.shannon , replicate.shannon_max = biodiversity_calculator.ShannonIndex( probs )
        replicate.simpson = biodiversity_calculator.SimpsonsIndex( counts, total )
        replicate.jack = biodiversity_calculator.jackKnife( counts, total )
        return replicate

    def run(self):
        """Loads the results in results_dict and populates a dictionary,
        with x-values as the keys, and y-values as the values.
        This allows multiple y-values to be stored per x-value.
        Optional queue will be used to send the results instead of
        returning them
        """
        results = Results()
        for number, nodes in self.rarefy( ):
            n_found_OTUs = len( nodes ) 
            if self.indices is not None:
                replicate = self._calc_indices( number , nodes ) 
                replicate.rarefaction = n_found_OTUs
            else:
                replicate = Replicate()
                replicate.nseqs = number
                replicate.rarefaction = n_found_OTUs
            results.append( replicate )
        if self.semaphore is not None:
            self.semaphore.release()
        if self.queue is not None:
            self.queue.put( results )
        return results 

    def rarefy( self ):
        """Given a results dictionary, this will sample random accessions
        from that dictionary and will yield the number of accessions
        sampled as well as the number of genera / species / OTUs that
        that a random sample of that number of accessions finds.
        """
        node = self.results_dict
        if self.options['-collapse-at-rank'] is not None:
            node = slib.collapse_at_rank( node , self.options , slib.TaxDB(), self.tax_index )
        accessions = self.accession_places.keys()
        n = self.nSeqs
        nsampled=self.increment
        multipled = True
        random_queue = Queue.Queue()
        random_generator = RandomNumberThread( random_queue , n , self.replicates, self.increment )
        random_generator.start()
        
        while nsampled < n:
            for replicate in xrange( self.replicates ):
                running_accessions = accessions[:]
                acc_places = self.accession_places.copy()
                n_left = n
                found_nodes = {}
                for ii in xrange( nsampled ):
                    accession = running_accessions.pop( random_queue.get() )
                    found_node = acc_places.pop( accession )
                    if found_node in found_nodes:
                        found_nodes[found_node].append( accession )
                    else:
                        found_nodes.update( { found_node : [accession] } )
                    n_left -= 1   # Every time we choose a random accession, the list of accessions is 1 shorter.
                    random_queue.task_done()
                yield nsampled , found_nodes
            nsampled += self.increment
        assert random_queue.get() == 'done'
        random_queue.task_done()
        random_generator.join()

class RandomNumberThread( threading.Thread ):
    def __init__(self,Q,N,replicates,increment):
        """Need a Queue.Queue object, number of accessions
        which we're rarefying, the number of replicates for
        each point, and the increment between points.
        """
        threading.Thread.__init__(self)
        self.N = N
        self.Q = Q
        self.increment = increment
        self.replicates = replicates
    def run( self ):
        i = self.increment
        while i < self.N:
            for replicate in xrange( self.replicates ):
                n_left = self.N
                for ii in xrange( i ):
                    rand = random.randint( 0, n_left-1 )
                    self.Q.put_nowait( rand )
                    n_left -= 1
            i += self.increment
        self.Q.put( 'done' )
        self.Q.join()
        return

def calc_increment( n_sequences , steps=None , calc_max_function=max ):
    """Calculate decent (x-axis) increment automatically.
    Will by default figure out the increment widths for 20 steps, for
    a number of sequences determined by calc_max_function.

    n_sequences should be a sequence containing the number of sequences
    for each set of results being plotted.

    calc_max_function is a function reference that will be called 
    with n_sequences, and the returned maximum will be divided by 
    `steps` (default: 20).

    e.g.
    >>> import numpy as np
    >>> # lambda function to calculate 95% confidence limits.
    >>> my_maximum = lambda x : np.average( x ) +  ( 1.96 * np.std( x ) )
    >>> increment = calc_increment( Loaders , steps=40 , calc_max_function=my_maximum )
    """
    if steps is None:
        steps = 20
    seqs = []
    maximum = calc_max_function( n_sequences )
    increment = maximum / steps
    return increment

def truncate_name( file_name ):
    # Might not be unique, but looks better; user specifies the order anyway.
    path_seps =  r'([/\\])'
    file_name = file_name.rsplit('.',1)[0]
    match = re.search( path_seps , file_name )
    if match:
        file_name = '/'.join( file_name.rsplit( match.group() , 2 )[-2:] )
    return file_name

def main(myoptions):
    artists = []
    if not matplot:
        save = {}

    plotter = PlotRarefaction( myoptions ) 

    # Draw pre-loaded results and exit.
    if myoptions['--draw']:
        all_results = {}
        for saved_file in myoptions['-in']:
            with file( saved_file ,'rb' ) as infile:
                results = pickle.load( infile )
            all_results.update( results['results'] )
        oldoptions = results['options']

        print myoptions['--indices'], myoptions['-indices']
        print oldoptions['--indices'], oldoptions['-indices']
        # Any other options need carrying over??
        oldoptions.update( { '-indices' : myoptions['-indices'] ,
                             '--indices': myoptions['--indices']} )
        plotter.options = oldoptions 
        plotter.resume( all_results )
        return
    
    separate_processes = {}
    semaphore = multiprocessing.Semaphore( multiprocessing.cpu_count())

    # Load data. Don't start threads yet. Got some counting to
    # do first.
    for infile in myoptions['-in']:
        with file( infile,'rb' ) as in_handle:
            data_dict = pickle.load( in_handle )

        data_name = truncate_name( infile )

        print 'Loading {0}'.format( data_name )
        separate_queue = multiprocessing.Queue()

        plotter.dataset_names.append( data_name )
        separate_process = Loader( data_dict, myoptions,
                queue=separate_queue, 
                semaphore = semaphore 
                )

        separate_processes.update( { data_name  :  (separate_process, separate_queue,) } )

    if myoptions['-increment'] is None:
        print 'Pre-analysing result files to determine a universal increment (go for 20 steps).'
        increment = calc_increment( proc[0].nSeqs for proc in separate_processes.values() )
        plotter.increment = increment
        print 'Sampling every {0} sequences.'.format( increment )

    # start the processes!!
    for process in separate_processes.itervalues():
        semaphore.acquire()
        process[0].increment = increment
        process[0].start()

    # Collate results and close processes.
    n_done = 0
    save = { 'options' : myoptions.options ,
             'results' : {} ,
           }
    processes_left = set( separate_processes.keys() )
    while n_done < len( separate_processes ):
        for data_name in list( processes_left ):
            process, queue = separate_processes[data_name]
            try:
                data = queue.get( timeout=3 )
                print 'Rarefied {0}.'.format( data_name )
                n_done += 1
            except (Queue.Empty, IOError):
                continue
            processes_left.remove( data_name )
            data.file_name = data_name
            process.join()
            if matplot:
                plotter.load_rarefaction( data )
            save['results'].update( { data_name : data } )

    # If we can, plot.
    if matplot:
        plotter.draw( )

    # Save the results to plot later if you want...
    if myoptions['-out'] == []:
        myoptions['-out'] = [ get_save_name() ]
    if myoptions['-out'] != ([] or [None]):
        with file( os.path.realpath( myoptions['-out'][0] ), 'wb' ) as outfile:
            pickle.dump( save , outfile , -1 )

def get_save_name():
    """Asks a user for a file name. If it exists, ask again,
    if user types nothing, ask again...
    Return None if user types 'n' or 'no'.
    """
    save_name = raw_input( 'Please enter a name to save the results, which can be drawn later:-\n')
    if save_name == ( 'n' or 'no' ):
        return None
    elif os.path.exists( save_name ):
        print 'That already exists. Trying again... (`no` or `n` to skip saving)'
        return get_save_name()
    elif save_name == '':
        print "If you wanna leave without saving, say 'no'!"
        return get_save_name()
    else:
        return save_name

if __name__ == '__main__':
    args = sys.argv[1:]
    myoptions = LocalOptions(args)
    main(myoptions)
else:
    myoptions = LocalOptions(  )
